/*-------------------------------------------------------------------------
 *
 * pg_statistics.c
 *		  用于操作统计数据表和修改语句表，以及启动清理进程
 * * Copyright (c) 2024, ajian
 *
 *-------------------------------------------------------------------------
 */

#include "pg_school.h"

#include "postgres.h"
#include "executor/executor.h"
#include "commands/async.h"
#include "tcop/utility.h"
#include "utils/elog.h"
#include "fmgr.h"
#include "miscadmin.h"
#include "executor/spi.h"
#include "utils/memutils.h"
#include "storage/ipc.h"
#include "postmaster/bgworker.h"
#include "utils/guc.h"
#include "commands/dbcommands.h"
#include "access/xact.h"      
#include "utils/snapmgr.h"   

void _PG_init(void);
void _PG_fini(void);
static void SignalHandlerForShutdown(SIGNAL_ARGS);
void CleanupWorkerMain(Datum main_arg);
void get_operateType(int op, char* operateType);
/* GUC 参数定义 */
static int cleanup_threshold = 100;
static int cleanup_interval = 60;
char *SQLlog_score = NULL;//记录score表修改语句表的插入
char *SQLlog_student = NULL;//记录student表修改语句表的插入
char *ROWSlog_score = NULL;//记录score表统计数据表的插入
char *ROWSlog_student = NULL;//记录student表统计数据表的插入
static ExecutorFinish_hook_type prev_ExecutorFinish_hook = NULL;

/* 枚举需要进行记录的一些函数调用操作 */
typedef enum {
    FUNC_INSERT_STUDENTS_NOSPI,
    FUNC_DELETE_STUDENTS_NOSPI,
    FUNC_QUERY_STUDENTS_NOSPI,
    FUNC_UPDATE_STUDENTS_NOSPI,
    FUNC_INSERT_STUDENT_SPI,
    FUNC_UPDATE_STUDENT_SPI,
    FUNC_DELETE_STUDENT_SPI,
    FUNC_QUERY_STUDENT_SPI,
    FUNC_INSERT_SCORE_NOSPI,
    FUNC_DELETE_SCORES_NOSPI,
    FUNC_QUERY_SCORES_NOSPI,
    FUNC_UPDATE_SCORES_NOSPI,
    FUNC_INSERT_SCORE_SPI,
    FUNC_UPDATE_SCORE_SPI,
    FUNC_DELETE_SCORE_SPI,
    FUNC_QUERY_SCORE_SPI,
    FUNC_UNKNOWN
} FunctionType;

/* 函数名数组与函数类型对应的映射 */
const char *function_names[] = {
    "insert_students_nospi",
    "delete_students_nospi",
    "query_students_nospi",
    "update_students_nospi",
    "insert_student_spi",
    "update_student_spi",
    "delete_student_spi",
    "query_student_spi",
    "insert_score_nospi",
    "delete_scores_nospi",
    "query_scores_nospi",
    "update_scores_nospi",
    "insert_score_spi",
    "update_score_spi",
    "delete_score_spi",
    "query_score_spi",
    NULL
};

/* 函数类型和操作类型的映射 */
OperationType function_to_operation_map[] = {
    OP_INSERT,  /* insert_students_nospi */
    OP_DELETE,  /* delete_students_nospi */
    OP_QUERY, /* query_students_nospi */
    OP_UPDATE,  /* update_students_nospi */
    OP_INSERT,  /* insert_student_spi */
    OP_UPDATE,  /* update_student_spi */
    OP_DELETE,  /* delete_student_spi */
    OP_QUERY, /* query_student_spi */
    OP_INSERT,  /* insert_score_nospi */
    OP_DELETE,  /* delete_scores_nospi */
    OP_QUERY, /* query_scores_nospi */
    OP_UPDATE,  /* update_scores_nospi */
    OP_INSERT,  /* insert_score_spi */
    OP_UPDATE,  /* update_score_spi */
    OP_DELETE,  /* delete_score_spi */
    OP_QUERY  /* query_score_spi */
};

/* 获取当前用户 */
static char *get_current_user() {
    Oid user_id = GetUserId();
    return pstrdup(GetUserNameFromId(user_id, false));
}



void get_operateType(int op, char* operateType) {
    switch (op) {
        case 1:
            strcpy(operateType, "insert_rows");
            break;
        case 2:
            strcpy(operateType, "delete_rows");
            break;
        case 3:
            strcpy(operateType, "update_rows");
            break;
        default:
            elog(WARNING, "Invalid operation code: %d", op);
            return; 
    }
}

/* 检查是否需要记录操作 */
static bool should_log_rows(const char *guc_setting, int op) {
    if (!guc_setting) {
        return false;
    }

    switch (op) {
        case insertType:
            return strstr(guc_setting, "insert") != NULL;
        case updateType:
            return strstr(guc_setting, "update") != NULL;
        case deleteType:
            return strstr(guc_setting, "delete") != NULL;
        default:
            return false;
    }
}


/*
*   往 operation_rows_log 表中记录修改语句影响的行数
*   根据自定义 GUC 参数来决定是否执行记录
*/
int log_operateRows(char *tablename, int op, int rows) {
    char *query = NULL;
    char operate[20]; // 定义操作类型的缓冲区
    int ret = 0;
    bool should_log = false;// 判断是否记录行数
    // 获取操作类型名称
    get_operateType(op, operate);
    // 根据表名决定是否记录行数
    if (strcmp(tablename, "score") == 0 && should_log_rows(ROWSlog_score,op)) {
        should_log = true;
    } else if (strcmp(tablename, "student") == 0 && should_log_rows(ROWSlog_student,op)) {
        should_log = true;
    }

    // 如果 should_log 为 true，才记录行数
    if (should_log) {
        query = psprintf("UPDATE %s SET %s = %s + %d WHERE table_name = '%s'",
                         SCHEMATABLEROWS, operate, operate, rows, tablename);

        SPI_connect();
        elog(LOG, "RECORD MODIFIED ROWS: %s", query);
        ret = SPI_exec(query, 0);
        SPI_finish();
        pfree(query);
    } else {
        elog(LOG, "Logging for operation on table '%s' is disabled by GUC parameters.", tablename);
    }

    return ret;
}

PG_FUNCTION_INFO_V1(reset_log_table);

/*
* 根据不同的表(统计表和修改语句表)和操作类型对数据进行置0或删除
*/
Datum reset_log_table(PG_FUNCTION_ARGS)
{
    char *table_name = text_to_cstring(PG_GETARG_TEXT_PP(0));   // 获取表名
    char *operation_type = text_to_cstring(PG_GETARG_TEXT_PP(1));  // 获取操作类型
    char *query = NULL;
    int ret = 0;
    int rows = 0;

    // 处理 operation_rows_log 表
    if (strcmp(table_name, TABLEROWS) == 0)
    {
        char *target_table = text_to_cstring(PG_GETARG_TEXT_PP(2));  // 获取目标表名 (score 或 student)

        if (strcmp(operation_type, "insert") == 0)
        {
            query = psprintf("UPDATE %s SET insert_rows = 0 WHERE table_name = '%s'", SCHEMATABLEROWS, target_table);
        }
        else if (strcmp(operation_type, "update") == 0)
        {
            query = psprintf("UPDATE %s SET update_rows = 0 WHERE table_name = '%s'", SCHEMATABLEROWS, target_table);
        }
        else if (strcmp(operation_type, "delete") == 0)
        {
            query = psprintf("UPDATE %s SET delete_rows = 0 WHERE table_name = '%s'", SCHEMATABLEROWS, target_table);
        }
        else
        {
            elog(WARNING, "Invalid operation type: %s", operation_type);
            PG_RETURN_NULL();
        }
    }
    // 处理 TABLESQL 表
    else if (strcmp(table_name, TABLESQL) == 0)
    {
        char *log_time = text_to_cstring(PG_GETARG_TEXT_PP(2));  // 获取时间戳

        query = psprintf("DELETE FROM %s WHERE log_time < '%s'", SCHEMATABLESQL, log_time);
        elog(LOG,"EXECUTE RESET SQL:%s",query);
    }
    else
    {
        elog(ERROR, "Invalid table name: %s", table_name);
        PG_RETURN_NULL();
    }

    // 连接 SPI
    SPI_connect();

    // 执行查询
    ret = SPI_exec(query, 0);
    if (ret != SPI_OK_UPDATE && ret != SPI_OK_DELETE)
    {
        elog(ERROR, "Query execution failed: %s", query);
    }

    rows = SPI_processed;

    // 清理
    SPI_finish();
    pfree(query);

    PG_RETURN_INT32(rows);
}


/* 对需要保存的内容进行转义 */
static char *escape_query(const char *query) {
    int len = strlen(query);
    int escaped_len = 0;
    char *escaped_query = NULL;
    int j = 0;
    for (int i = 0; i < len; i++) {
        if (query[i] == '\'') {
            escaped_len++;
        }
    }
    escaped_query = palloc(len + escaped_len + 1);

    for (int i = 0; i < len; i++) {
        if (query[i] == '\'') {
            escaped_query[j++] = '\\';
        }
        escaped_query[j++] = query[i];
    }
    escaped_query[j] = '\0';
    return escaped_query;
}

static char *truncate_query(const char *query) {

    int len = strlen(query);
    int mask_len;
    int remaining_length; // 保留前后各一半
    int prefix_length;
    int suffix_length;
    char *truncated_query = NULL;
    const char *mask = "******";
    if (len <= 1024) {
        return pstrdup(query);
    }
    // 计算替换后的长度
    
    mask_len = strlen(mask);
    remaining_length = 1024 - (mask_len * 2); // 保留前后各一半
    prefix_length = remaining_length / 2;
    suffix_length = remaining_length - prefix_length;

    truncated_query = palloc(1025); // 为 '\0' 留出空间
    strncpy(truncated_query, query, prefix_length);
    strcpy(truncated_query + prefix_length, mask);
    strcpy(truncated_query + prefix_length + mask_len, query + len - suffix_length);
    truncated_query[1024] = '\0'; // 确保字符串以 \0 结束

    return truncated_query;
}
/*
*   记录调用的SQL语句到修改语句表中
*/
static void log_operateSQL(const char *query_text) {
    char *escaped_query = escape_query(query_text);
    char *truncated_query = truncate_query(escaped_query);
    char *current_user = get_current_user();
    char *query = NULL;
    int ret;

    query = psprintf("INSERT INTO %s (query_text, modified_by) VALUES (E'%s', '%s');", \
                                        SCHEMATABLESQL, truncated_query, current_user);
    
    SPI_connect();
    ret = SPI_exec(query, 0);
    if (ret != SPI_OK_INSERT) {
        elog(ERROR, "Failed to log query: %s", query);
    }
    SPI_finish();
    
    pfree(escaped_query);
    pfree(truncated_query);
    pfree(current_user);
    pfree(query);
}

/* 获取函数的类型 */
static FunctionType get_function_type(Oid funcid) {
    char *funcname = get_func_name(funcid);
    if (!funcname) {
        return FUNC_UNKNOWN;
    }

    for (int i = 0; function_names[i] != NULL; i++) {
        if (strcmp(funcname, function_names[i]) == 0) {
            pfree(funcname);  // 释放内存
            return (FunctionType) i;  // 返回对应的枚举值
        }
    }

    pfree(funcname);  // 释放内存
    return FUNC_UNKNOWN;
}

/* 记录日志的辅助函数 */
static void log_operation(OperationType op_type, const char *schemaname, \
                                const char *relname, const char *query_string) {
    const char *op_str = (op_type == OP_INSERT) ? "INSERT" :
                         (op_type == OP_UPDATE) ? "UPDATE" :
                         (op_type == OP_DELETE) ? "DELETE" : "UNKNOWN";
    elog(LOG, "Captured operation: %s on table: %s.%s", op_str, schemaname, relname);
    elog(LOG, "Captured query: %s", query_string);
    log_operateSQL(query_string);  // 执行日志记录
}

/* 检查是否需要记录操作 */
static bool should_log_operation(const char *guc_setting, OperationType op_type) {
    if (!guc_setting) {
        return false;
    }
    switch (op_type) {
        case OP_INSERT:
            return strstr(guc_setting, "insert") != NULL;
        case OP_UPDATE:
            return strstr(guc_setting, "update") != NULL;
        case OP_DELETE:
            return strstr(guc_setting, "delete") != NULL;
        default:
            return false;
    }
}

/* ExecutorFinish Hook 实现 */
static void my_ExecutorFinish(QueryDesc *queryDesc) {
    const char *query_string = queryDesc->sourceText;
    bool should_log = false;
    OperationType op_type = OP_QUERY;
    ListCell *lc = NULL;
    int rows = 0;

    PushActiveSnapshot(GetTransactionSnapshot());

    /* 检查是否有 UDF 函数调用 */
    foreach(lc, queryDesc->plannedstmt->planTree->targetlist) {
        TargetEntry *tle = (TargetEntry *) lfirst(lc);
        if (IsA(tle->expr, FuncExpr)) {
            FuncExpr *funcexpr = (FuncExpr *) tle->expr;
            Oid funcid = funcexpr->funcid;

            /* 获取函数类型并查找相应操作类型 */
            FunctionType func_type = get_function_type(funcid);
            if (func_type != FUNC_UNKNOWN) {
                OperationType udf_op_type = function_to_operation_map[func_type];
                const char *relname = (func_type < 8) ? "student" : "score";
                if ((strcmp(relname, TABLESCORE) == 0 && should_log_operation(SQLlog_score, udf_op_type)) ||
                    (strcmp(relname, TABLESTUDENT) == 0 && should_log_operation(SQLlog_student, udf_op_type))) {
                    log_operation(udf_op_type, SCHEMANAME, relname, query_string);
                    should_log = true;
                }
            }
        }
    }


    // 如果没有调用 UDF 函数，则判断用户的原生增删改操作
    if (!should_log && !is_udf_execution) {
        /* 判断是普通表操作 (INSERT, UPDATE, DELETE) */
        if (queryDesc->operation == CMD_INSERT) {
            op_type = OP_INSERT;
        } else if (queryDesc->operation == CMD_UPDATE) {
            op_type = OP_UPDATE;
        } else if (queryDesc->operation == CMD_DELETE) {
            op_type = OP_DELETE;
        }

        /* 遍历 RangeTblEntry 以查看目标表 */
        foreach(lc, queryDesc->estate->es_range_table) {
            RangeTblEntry *rte = (RangeTblEntry *) lfirst(lc);
            if (rte->rtekind == RTE_RELATION) {
                /* 获取表名和模式名 */
                char *relname = get_rel_name(rte->relid);
                char *schemaname = get_namespace_name(get_rel_namespace(rte->relid));

                if (relname && schemaname && strcmp(schemaname, SCHEMANAME) == 0) {
                    /* 只处理 TABLESCORE 和 TABLESTUDENT 表 */
                    if (strcmp(relname, TABLESCORE) == 0) {
                        //是否需要记录修改语句表指定操作
                        if(should_log_operation(SQLlog_score, op_type)){
                            log_operation(op_type, schemaname, relname, query_string);
                            should_log = true;
                        }
                        //获取实际影响行数
                        rows = queryDesc->estate->es_processed;
                        //记录操作行数
                        log_operateRows(TABLESCORE, op_type, rows);
                    } else if (strcmp(relname, TABLESTUDENT) == 0) {
                        //是否需要记录修改语句表指定操作
                        if(should_log_operation(SQLlog_student, op_type))
                        {
                            log_operation(op_type, schemaname, relname, query_string);
                            should_log = true;
                        }
                        //获取实际影响行数
                        rows = queryDesc->estate->es_processed;
                        //记录操作行数
                        log_operateRows(TABLESTUDENT, op_type, rows);
                    }

                }

                if (relname) pfree(relname);
                if (schemaname) pfree(schemaname);
            }
        }
    }

    PopActiveSnapshot();

    /* 调用上一个钩子或标准 ExecutorFinish */
    if (prev_ExecutorFinish_hook) {
        prev_ExecutorFinish_hook(queryDesc);
    } else {
        standard_ExecutorFinish(queryDesc);
    }
}


/* Background worker main function */
void CleanupWorkerMain(Datum main_arg) {
    BackgroundWorkerUnblockSignals();
    pqsignal(SIGTERM, SignalHandlerForShutdown);

    /* Connect to the specified database */
    BackgroundWorkerInitializeConnection("postgres", NULL, 0);
    elog(LOG, "pgschool_cleanup_worker has started successfully.");

    while (true) {
        int ret = 0;
        int rows = 0;
        Snapshot snapshot = NULL;
        char *query = psprintf("DELETE FROM %s WHERE log_time IN (" 
                            "SELECT log_time FROM %s "
                            "ORDER BY log_time DESC "
                            "OFFSET %d)", 
                            SCHEMATABLESQL, SCHEMATABLESQL, cleanup_threshold);

        elog(LOG, "EXECUTE CLEANUP SQL: %s", query);
        /* 开始一个事务 */
        StartTransactionCommand();
        /* 创建快照 */
        snapshot = GetTransactionSnapshot();
        PushActiveSnapshot(snapshot);
        /* 连接 SPI */
        ret = SPI_connect();
        if (ret != SPI_OK_CONNECT) {
            PopActiveSnapshot();
            AbortCurrentTransaction();
            pfree(query);
            pg_usleep(cleanup_interval * 1000000L);
            continue;
        }
        /* 执行清理查询 */
        ret = SPI_exec(query, 0);
        if (ret != SPI_OK_DELETE) {
            elog(WARNING, "CleanupWorker: Failed to delete old records");
            SPI_finish();
            PopActiveSnapshot();
            AbortCurrentTransaction();
            pfree(query);
            pg_usleep(cleanup_interval * 1000000L);
            continue;
        }
        rows = SPI_processed;
        //清理进程是否清理实际记录，通过日志告知
        if(rows > 0)
        {
             elog(LOG, "CleanupWorker work success: clean %d record from %s",rows,TABLESQL);
        }
        if(rows == 0)
        {
             elog(LOG, "CleanupWorker work not effort: clean %d record from %s",rows, TABLESQL);
        }        
        SPI_finish();

        /* 清理快照并提交事务 */
        PopActiveSnapshot();
        CommitTransactionCommand();

        pfree(query);

        /* Sleep until next interval */
        pg_usleep(cleanup_interval * 1000000L);
    }
}

static void SignalHandlerForShutdown(SIGNAL_ARGS) {
    proc_exit(0);
}

/* Register the background worker */
void _PG_init(void) {
    BackgroundWorker worker;
    //如果不在posgres.conf中配置shared_preload_libraries，报错
    if (!process_shared_preload_libraries_in_progress)
    {
        elog(ERROR,"pg_school can only be loaded via shared_preload_libraries");
        return;
    }

    /* Define custom GUC parameters */
    DefineCustomIntVariable("changjunhighschool.cleanup_threshold",
                            "Number of rows after which the table should be cleaned up.",
                            NULL,
                            &cleanup_threshold,
                            1000,
                            1,
                            INT_MAX,
                            PGC_SIGHUP,
                            0,
                            NULL,
                            NULL,
                            NULL);

    DefineCustomIntVariable("changjunhighschool.cleanup_interval",
                            "Interval (in seconds) between cleanup runs.",
                            NULL,
                            &cleanup_interval,
                            60,
                            1,
                            INT_MAX,
                            PGC_SIGHUP,
                            0,
                            NULL,
                            NULL,
                            NULL);

    // Define custom GUC parameters for logging configuration
    DefineCustomStringVariable("changjunhighschool.SQLlog_score",
                               "Comma-separated list of operations to log for the score table (insert, update, delete).",
                               NULL,
                               &SQLlog_score, 
                               "insert,update,delete", 
                               PGC_SIGHUP,
                               0,
                               NULL,
                               NULL,
                               NULL);

    /* Define GUC parameter for logging student table operations */
    DefineCustomStringVariable("changjunhighschool.SQLlog_student",
                               "Comma-separated list of operations to log for the student table (insert, update, delete).",
                               NULL,
                               &SQLlog_student, 
                               "insert,update,delete", 
                               PGC_SIGHUP,
                               0,
                               NULL,
                               NULL,
                               NULL);

    DefineCustomStringVariable("changjunhighschool.ROWSlog_student",
                               "Comma-separated list of operations to log for the student table (insert, update, delete).",
                               NULL,
                               &ROWSlog_student, 
                               "insert,update,delete", 
                               PGC_SIGHUP,
                               0,
                               NULL,
                               NULL,
                               NULL);
    DefineCustomStringVariable("changjunhighschool.ROWSlog_score",
                               "Comma-separated list of operations to log for the student table (insert, update, delete).",
                               NULL,
                               &ROWSlog_score, 
                               "insert,update,delete", 
                               PGC_SIGHUP,
                               0,
                               NULL,
                               NULL,
                               NULL);

    /* Register the background worker */
    memset(&worker, 0, sizeof(BackgroundWorker));
    worker.bgw_flags = BGWORKER_SHMEM_ACCESS | BGWORKER_BACKEND_DATABASE_CONNECTION;
    worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
    worker.bgw_restart_time = cleanup_interval;
    snprintf(worker.bgw_name, BGW_MAXLEN, "pgschool_cleanup_worker");
    snprintf(worker.bgw_library_name, BGW_MAXLEN, "pg_school");
    snprintf(worker.bgw_function_name, BGW_MAXLEN, "CleanupWorkerMain");
    worker.bgw_main_arg = (Datum) 0;
    worker.bgw_notify_pid = 0;

    RegisterBackgroundWorker(&worker);

    /* 注册 ExecutorFinish 钩子 */
    prev_ExecutorFinish_hook = ExecutorFinish_hook;
    ExecutorFinish_hook = my_ExecutorFinish;
}

/* 模块卸载函数 */
void _PG_fini(void) {
    elog(INFO,"pgschool_cleanup_worker process exit by DROPING EXTENSION");
    ExecutorFinish_hook = prev_ExecutorFinish_hook;
}